/*
使用百度的ocr识别
*/
package main

import (
	"bytes"
	"crypto/tls"
	"encoding/base64"
	"errors"
	"fmt"
	"github.com/suiyunonghen/DxCommonLib"
	"github.com/suiyunonghen/dxsvalue"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

var DefHttpClient = http.Client{Transport: &http.Transport{
	TLSClientConfig:       &tls.Config{InsecureSkipVerify: true}, //跳过证书的有效认证
	MaxIdleConns:          100,
	MaxIdleConnsPerHost:   2,
	IdleConnTimeout:       90 * time.Second,
	TLSHandshakeTimeout:   10 * time.Second,
	ExpectContinueTimeout: 1 * time.Second,
	DialContext: (&net.Dialer{
		Timeout:   30 * time.Second,
		KeepAlive: 30 * time.Second,
	}).DialContext,
},
	Timeout: time.Minute * 5, //设置超时时间，5分钟
}

type Token struct {
	shareBinary   [4096]byte
	RefreshToken  string
	ExpiresIn     int
	SessionKey    string
	AccessToken   string
	Scope         string
	SessionSecret string
}

type Location struct {
	Width  int
	Top    int
	Left   int
	Height int
}

type WordsResult struct {
	Location Location `json:"location"`
	Words    string   `json:"words"`
}

type OcrResult struct {
	shareBinary    [4096]byte
	LogId          uint64
	WordsResultNum int
	WordsResult    []WordsResult
	ErrorCode      int
	ErrorMsg       string
}

func (result *OcrResult) ReadFromJsonBytes(bt []byte) error {
	value := getValue()
	err := value.LoadFromJson(bt, true)
	if err != nil {
		freeValue(value)
		return err
	}
	result.ErrorMsg = value.AsString("error_msg", "")
	result.ErrorCode = value.AsInt("error_code", 0)
	result.LogId = uint64(value.AsInt64("log_id", 0))
	result.WordsResultNum = value.AsInt("words_result_num", 0)
	item := value.ValueByName("words_result")
	if item != nil {
		for i := 0; i < item.Count(); i++ {
			_, vitem := item.Items(i)
			var wresult WordsResult
			wresult.Words = vitem.AsString("words", "")
			locationItem := vitem.ValueByName("location")
			if locationItem != nil {
				wresult.Location.Width = locationItem.AsInt("width", 0)
				wresult.Location.Height = locationItem.AsInt("height", 0)
				wresult.Location.Left = locationItem.AsInt("left", 0)
				wresult.Location.Top = locationItem.AsInt("top", 0)
			}
			result.WordsResult = append(result.WordsResult, wresult)
		}
	}
	freeValue(value)
	return nil
}

func (result *OcrResult) getLeftMiddle() int {
	var max, min int
	if len(result.WordsResult) == 0 {
		return 0
	} else {
		min = result.WordsResult[0].Location.Left
		max = min
	}
	for _, v := range result.WordsResult {
		if min > v.Location.Left {
			min = v.Location.Left
		}
		if max < v.Location.Left {
			max = v.Location.Left
		}
	}
	return (max + min) / 2
}

func (result *OcrResult) getDistMiddle() int {
	var max, min int
	var last int
	if len(result.WordsResult) < 2 {
		return 0
	} else {
		last = result.WordsResult[1].Location.Top
		min = last - result.WordsResult[0].Location.Top
		max = min
	}
	for i := 2; i < len(result.WordsResult); i++ {
		dist := result.WordsResult[i].Location.Top - last
		last = result.WordsResult[i].Location.Top
		if min > dist {
			min = dist
		}
		if max < dist {
			max = dist
		}
	}
	if (max - min) > 10 {
		return (max + min) / 2
	} else {
		return max
	}
}

func (result *OcrResult) OcrText() string {
	var last, lastH int
	var out strings.Builder
	if len(result.WordsResult) > 0 {
		last = result.WordsResult[0].Location.Top
		lastH = result.WordsResult[0].Location.Height

		out.WriteString(result.WordsResult[0].Words)
	}
	for i := 1; i < len(result.WordsResult); i++ {
		d := result.WordsResult[i].Location.Top - last
		last = result.WordsResult[i].Location.Top
		if float32(d) >= float32(lastH)/2 {
			out.WriteString("\n" + result.WordsResult[i].Words)
		} else {
			out.WriteString(result.WordsResult[i].Words)
		}
		lastH = result.WordsResult[i].Location.Height
	}
	return out.String()
	/*dm := result.getDistMiddle()
	lm := result.getLeftMiddle()
	var out strings.Builder
	var last int
	if len(result.WordsResult) > 0 {
		last = result.WordsResult[0].Location.Top
		out.WriteString(result.WordsResult[0].Words)
	}
	//组织判断是否需要换行
	for i := 1; i < len(result.WordsResult); i++ {
		d := result.WordsResult[i].Location.Top - last
		last = result.WordsResult[i].Location.Top
		if d > dm {
			out.WriteString("\n" + result.WordsResult[i].Words)
		} else if result.WordsResult[i].Location.Left <= lm {
			out.WriteString(result.WordsResult[i].Words)
		} else {
			out.WriteString("\n" + result.WordsResult[i].Words)
		}
	}
	//out.WriteString("\n")
	return out.StatusString()*/
}

var (
	resultPool sync.Pool
	valuePool  sync.Pool
)

func getResult() *OcrResult {
	v := resultPool.Get()
	if v == nil {
		result := new(OcrResult)
		result.WordsResult = make([]WordsResult, 0, 32)
		return result
	} else {
		result := v.(*OcrResult)
		return result
	}
}

func getValue() *dxsvalue.DxValue {
	v := valuePool.Get()
	if v == nil {
		result := dxsvalue.NewObject(true)
		return result
	} else {
		return v.(*dxsvalue.DxValue)
	}
}

func freeValue(v *dxsvalue.DxValue) {
	v.Clear()
	valuePool.Put(v)
}

func FreeResult(r *OcrResult) {
	r.WordsResult = r.WordsResult[:0]
	resultPool.Put(r)
}

type BaiduOcrClient struct {
	appID            string
	apiKEY           string
	secretKEY        string
	fActive          int32
	fAccurateCount   int32 //当天精准识别个数
	fAccurateOutTime atomic.Value
	FGenCount        int32 //当天普通识别个数
	fGenOutTime      atomic.Value
	lastTokenTime    atomic.Value
	ftoken           chan struct{} //百度一次只能2QPS并发，做一个令牌
	token            Token
}

func (client *BaiduOcrClient) getToken() (err error) {
	/*获取Access Token
	请求URL数据格式
	向授权服务地址https://aip.baidubce.com/oauth/2.0/token发送请求（推荐使用POST），并在URL中带上以下参数：
	grant_type： 必须参数，固定为client_credentials；
	client_id： 必须参数，应用的API Key；
	client_secret： 必须参数，应用的Secret Key；*/
	var tkUrl = "https://aip.baidubce.com/oauth/2.0/token"
	var grantType = "client_credentials"
	var reqUrl = fmt.Sprintf("%s?grant_type=%s&client_id=%s&client_secret=%s", tkUrl, grantType, client.apiKEY, client.secretKEY)
	req, err := http.NewRequest("GET", reqUrl, nil)
	if err != nil {
		return err
	}
	req.Header.Add("Content-Type", "application/json; charset=UTF-8")
	resp, err := DefHttpClient.Do(req)
	if err != nil {
		return err
	}
	/*bt,_ := ioutil.ReadAll(resp.Body)
	fmt.Println(len(bt))
	fmt.Println(string(bt))*/
	l, err := resp.Body.Read(client.token.shareBinary[:])
	resp.Body.Close()
	if err != nil && err != io.EOF {
		return err
	}

	v := getValue()
	err = v.LoadFromJson(client.token.shareBinary[:l], true)
	if err != nil {
		freeValue(v)
		return err
	}
	client.token.RefreshToken = v.StringByName("refresh_token", "")
	client.token.ExpiresIn = v.IntByName("expires_in", 0)
	client.token.AccessToken = v.StringByName("access_token", "")
	client.token.SessionKey = v.StringByName("session_key", "")
	client.token.Scope = v.StringByName("scope", "")
	client.token.SessionSecret = v.StringByName("session_secret", "")
	client.lastTokenTime.Store(time.Now())
	freeValue(v)
	atomic.StoreInt32(&client.fActive, 1)
	return nil
}

// Accurate 高精度识别
func (client *BaiduOcrClient) accurate(image []byte) (res *OcrResult, err error) {
	var acUrl = "https://aip.baidubce.com/rest/2.0/ocr/v1/accurate"
	var aUrl = fmt.Sprintf("%s?access_token=%s", acUrl, client.token.AccessToken)
	var data = bytes.NewBufferString("")
	v := url.Values{}
	imgEnc := make([]byte, base64.StdEncoding.EncodedLen(len(image)))
	base64.StdEncoding.Encode(imgEnc, image)
	v.Add("image", string(imgEnc))
	data.WriteString(v.Encode())

	client.ftoken <- struct{}{} //获取
	startTime := time.Now()
	defer func(ltime time.Time) {
		DxCommonLib.MustRunAsync(client.waitQps, ltime)
	}(startTime)
	for idx := 0; idx < 3; idx++ {
		resp, err := DefHttpClient.Post(aUrl, "application/x-www-form-urlencoded", data)
		if err != nil {
			return nil, err
		}
		if res == nil {
			res = getResult()
		}
		l, err := resp.Body.Read(res.shareBinary[:])
		resp.Body.Close()
		if err != nil && err != io.EOF {
			FreeResult(res)
			return nil, err
		}
		if err = res.ReadFromJsonBytes(res.shareBinary[:l]); err != nil {
			FreeResult(res)
			res = nil
		}
		if res.ErrorCode != 1 && res.ErrorCode != 2 {
			if res.ErrorCode == 17 {
				atomic.StoreInt32(&client.fAccurateCount, 500) //设定为已经到极限了
				client.fAccurateOutTime.Store(time.Now())
			} else if res.ErrorCode == 0 {
				atomic.AddInt32(&client.fAccurateCount, 1)
			}
			return res, nil
		}
	}

	return res, err
}

func (client *BaiduOcrClient) waitQps(data ...interface{}) {
	ltime := data[0].(time.Time)
	ctime := time.Now()
	t := ctime.Sub(ltime)
	var waitLen time.Duration
	if t < time.Second {
		waitLen = time.Second - t
	} else {
		waitLen = time.Millisecond * 100
	}
	DxCommonLib.Sleep(waitLen)
	<-client.ftoken //释放
}

// General 通用识别
func (client *BaiduOcrClient) general(image []byte) (res *OcrResult, err error) {
	var acUrl = "https://aip.baidubce.com/rest/2.0/ocr/v1/general"
	var aUrl = fmt.Sprintf("%s?access_token=%s", acUrl, client.token.AccessToken)
	var data = bytes.NewBufferString("")
	v := url.Values{}
	imgEnc := make([]byte, base64.StdEncoding.EncodedLen(len(image)))
	base64.StdEncoding.Encode(imgEnc, image)
	v.Add("image", string(imgEnc))
	data.WriteString(v.Encode())

	client.ftoken <- struct{}{} //获取
	startTime := time.Now()
	defer func(ltime time.Time) {
		DxCommonLib.MustRunAsync(client.waitQps, ltime)
	}(startTime)
	for idx := 0; idx < 3; idx++ {
		resp, err := DefHttpClient.Post(aUrl, "application/x-www-form-urlencoded", data)
		if err != nil {
			return nil, err
		}
		if res == nil {
			res = getResult()
		}
		l, err := resp.Body.Read(res.shareBinary[:])
		resp.Body.Close()
		if err != nil && err != io.EOF {
			FreeResult(res)
			return nil, err
		}
		if err = res.ReadFromJsonBytes(res.shareBinary[:l]); err != nil {
			FreeResult(res)
			res = nil
		}
		if res.ErrorCode != 1 && res.ErrorCode != 2 {
			if res.ErrorCode == 4 {
				atomic.StoreInt32(&client.FGenCount, 50000) //设定为已经到极限了
				client.fGenOutTime.Store(time.Now())
			} else if res.ErrorCode == 0 {
				atomic.AddInt32(&client.FGenCount, 1)
			}
			return res, nil
		}
	}
	return res, err
}

func (client *BaiduOcrClient) BestOcrPic(image []byte) (res *OcrResult, err error) {
	if atomic.LoadInt32(&client.fActive) == 0 {
		return nil, nil
	}
	if atomic.LoadInt32(&client.fAccurateCount) < 500 {
		res, err = client.accurate(image)
		if res.ErrorCode == 0 {
			return res, err
		}
	}
	return nil, nil
}

func (client *BaiduOcrClient) GenOcrPic(image []byte) (res *OcrResult, err error) {
	if atomic.LoadInt32(&client.fActive) == 0 {
		return nil, nil
	}
	if atomic.LoadInt32(&client.FGenCount) < 50000 {
		return client.general(image)

	}
	return nil, nil
}

func (client *BaiduOcrClient) OcrPic(image []byte) (res *OcrResult, err error) {
	if atomic.LoadInt32(&client.fActive) == 0 {
		return nil, nil
	}
	//百度免费版本，一次只能2个并发
	if atomic.LoadInt32(&client.fAccurateCount) < 500 {
		res, err = client.accurate(image)
		if res.ErrorCode == 0 {
			return res, err
		}
	}
	if atomic.LoadInt32(&client.FGenCount) < 50000 {
		return client.general(image)

	}
	return nil, nil
}

func (client *BaiduOcrClient) checkTokenAndTimeOut() {
	for {
		select {
		case <-DxCommonLib.After(time.Minute * 30):
			//30分钟检查一次
			curtime := time.Now()
			t := client.lastTokenTime.Load().(time.Time)
			if t.IsZero() || curtime.Sub(t) > time.Hour*24 { //超过一天，或者说没获取token的
				client.getToken()
			}
			//判定一下上次的时间
			if atomic.LoadInt32(&client.fAccurateCount) == 500 {
				atime := client.fAccurateOutTime.Load().(time.Time)
				if curtime.Day() != atime.Day() {
					client.fAccurateOutTime.Store(time.Now()) //又重新计算
					atomic.StoreInt32(&client.fAccurateCount, 0)
				}
			}
			if atomic.LoadInt32(&client.FGenCount) == 50000 {
				atime := client.fGenOutTime.Load().(time.Time)
				if curtime.Day() != atime.Day() {
					atomic.StoreInt32(&client.FGenCount, 0)
					client.fGenOutTime.Store(time.Now()) //又重新计算
				}
			}
		}
	}
}

func NewBaiduOcrClient(id, apiKey, secKey string) (res *BaiduOcrClient, err error) {
	result := BaiduOcrClient{appID: id, apiKEY: apiKey, secretKEY: secKey}
	result.lastTokenTime.Store(time.Time{})
	err = result.getToken()
	if err != nil {
		return nil, err
	}
	if result.token.AccessToken == "" {
		return nil, errors.New("Error get token")
	}
	result.ftoken = make(chan struct{}, 1)
	go result.checkTokenAndTimeOut() //刷新token以及检查是否有新的免费额度可以用
	return &result, nil
}
